from lib import *

def test_matmul():
    # 定义输入参数
    matrix_M = np.random.randint(1, 100)
    matrix_N = np.random.randint(1, 100)
    matrix_K = np.random.randint(1, 100)
    # 定义矩阵乘法参数
    layer = {
        'M': matrix_M,
        'N': matrix_N,
        'K': matrix_K
    }
    print(layer)

    # 创建输入数据
    left = np.random.randint(-128, 127, size=(matrix_M, matrix_K))
    right = np.random.randint(-128, 127, size=(matrix_K, matrix_N))
    bias = np.random.randint(-128, 127, size=(matrix_N))
    deq = np.random.rand(matrix_N)
    right_nk = right.transpose()
    result_matrix_mn = np.zeros((matrix_M, matrix_N))

    for l1_n_start_in_l2 in range(0, matrix_N, N_L1):
        n_size_l1 = min(N_L1, matrix_N - l1_n_start_in_l2)
        bias_n_l1 = bias[l1_n_start_in_l2:l1_n_start_in_l2+n_size_l1]
        deq_n_l1 = deq[l1_n_start_in_l2:l1_n_start_in_l2+n_size_l1]
        for l1_m_start_in_l2 in range(0, matrix_M, M_L1):
            m_size_l1 = min(M_L1, matrix_M - l1_m_start_in_l2)
            result_m2n2m1n1m0n0_psb = np.zeros((ceil_div(m_size_l1, M_L0), ceil_div(n_size_l1, N_L0), M1_L0, N1_L0, M0, N0))
            for l1_k_start_in_l2 in range(0, matrix_K, K_L1):
                k_size_l1 = min(K_L1, matrix_K - l1_k_start_in_l2)
                k_size_l1_align_k0 = (k_size_l1 + K0 - 1) // K0 * K0
                left_k1mk0_l1 = MK2K1MK0_DDR2L1(left, matrix_M, matrix_K, l1_m_start_in_l2, l1_k_start_in_l2, m_size_l1, k_size_l1)
                right_k1nk0_l1 = MK2K1MK0_DDR2L1(right_nk, matrix_N, matrix_K, l1_n_start_in_l2, l1_k_start_in_l2, n_size_l1, k_size_l1)
                for l0_n_start_in_l2 in range(0, n_size_l1, N_L0):
                    n_size_l0 = min(N_L0, n_size_l1 - l0_n_start_in_l2)
                    bias_n1n0_mmb = N2N1N0_L12MMB(bias_n_l1, l0_n_start_in_l2, n_size_l0)
                    deq_n1n0_mmb = N2N1N0_L12MMB(deq_n_l1, l0_n_start_in_l2, n_size_l0)
                    for l0_m_start_in_l2 in range(0, m_size_l1, M_L0):
                        m_size_l0 = min(M_L0, m_size_l1 - l0_m_start_in_l2)
                        for l0_k_start_in_l2 in range(0, k_size_l1_align_k0, K_L0):
                            k_size_l0 = min(K_L0, k_size_l1_align_k0 - l0_k_start_in_l2)
                            assert(k_size_l0 % K0 == 0)
                            k1_size_l0 = k_size_l0 // K0
                            l0_k1_start_in_l2 = l0_k_start_in_l2 // K0
                            left_m1k1m0k0_lmb = K1MK02M1K1M0K0_L12LMB(left_k1mk0_l1, l0_m_start_in_l2, l0_k1_start_in_l2, m_size_l0, k1_size_l0)
                            right_n1k1n0k0_rmb = K1NK02N1K1N0K0_L12RMB(right_k1nk0_l1, l0_n_start_in_l2, l0_k1_start_in_l2, n_size_l0, k1_size_l0)
                            bias_en = (l0_k_start_in_l2==0 and l1_k_start_in_l2==0)
                            psum_en = l0_k_start_in_l2 > 0 or l1_k_start_in_l2 > 0
                            deq_en = (l0_k_start_in_l2 // K_L0 == ceil_div(k_size_l1_align_k0, K_L0) - 1) and (l1_k_start_in_l2 // K_L1 == ceil_div(matrix_K, K_L1) - 1)
                            result_m2n2m1n1m0n0_psb[l0_m_start_in_l2//M_L0, l0_n_start_in_l2//N_L0] = matmul_m1k1m0k0_n1k1n0k0(\
                            result_m2n2m1n1m0n0_psb[l0_m_start_in_l2//M_L0, l0_n_start_in_l2//N_L0], left_m1k1m0k0_lmb, right_n1k1n0k0_rmb, bias_n1n0_mmb, deq_n1n0_mmb, left_m1k1m0k0_lmb.shape[0], right_n1k1n0k0_rmb.shape[0], left_m1k1m0k0_lmb.shape[1], bias_en, psum_en, deq_en)
                        M1N1M0N02MN_PSB2DDR(result_matrix_mn, result_m2n2m1n1m0n0_psb[l0_m_start_in_l2//M_L0, l0_n_start_in_l2//N_L0], l1_m_start_in_l2+l0_m_start_in_l2, l1_n_start_in_l2+l0_n_start_in_l2, m_size_l0, n_size_l0)

    golden_matrix_mn = matmul_mk_kn(left, right, bias, deq)
    diff = result_matrix_mn - golden_matrix_mn
    if(np.abs(diff).mean() < 1e-3):
        print("gemm pass")
    else:
        print("gemm fail")
    return True

test_num = 10
while(test_num > 0):
    if(test_matmul()):
        test_num -= 1